from __future__ import division
from operator import itemgetter, attrgetter
import os
import sys
import math
# Updated by Wei 2013/07/27 night

# Updated by Wei 2013/07/22 afternoon, configured for computing the connected postings

# Updated by Wei 2013/07/20 afternoon 
# deal with a relative large training file in the following folder: 
# /data3/obukai/workspace/web-search-engine-wei/polyIRIndexer/distributedWorkForFeatureGenerationPostingRankInDoc20130720

# updated by Wei 2013/07/15
# This program will have ONE main logic and two functions.
# (1) function1: preprocess the training example file and output a list of documents and compressed files needed
# to feed into the polyIRToolkit
# (2) after the polyIRToolkit has outputed her own output file similar with the name called: gov2DocumentWithTheirCompletedPostingSet_for_debug_ONLY
# (3) function2: postprocess includes the following: regarded the gov2DocumentWithTheirCompletedPostingSet_for_debug_ONLY as the input and compute the 
# feature correct posting_rank_in_doc and outputting a new training file with the new feature added

def preprocess():
    print "preprocess(...) function called."
    # key: trecID
    # value: freq (how many times this document has been touched/selected by the sampling method)
    trecIDDict = {}
    trecIDList = []
    # key: path
    # value:freq
    correspondingGov2CompressedFilesCompletedPathsList = []
    # inputFileName = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/Training_Set_2013_07_13_tfqAdded_postingRankInListAdded_tfcAdded.train"
    # inputFileName = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/MLRelated/Model15/distributedWorkForFeatureGenerationPostingRankInDoc20130721ForConnectedPostings/xae"
    inputFileName = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/MLRelated/Model17/TOPK_sorted_by_queryID/distributedWorkForFeatureGenerationPostingRankInDoc20130727/xaj"
    inputRelativeFileName = inputFileName.strip().split("/")[-1]
    inputFileHandler = open(inputFileName,"r")
    # Switch ON/OFF for headline
    # inputFileHandler.readline()
    
    for line in inputFileHandler.readlines():
        lineElements = line.strip().split(" ")
        trecID = lineElements[2]
        if trecID not in trecIDDict:
            trecIDDict[trecID] = 1
        else:
            trecIDDict[trecID] += 1
    inputFileHandler.close()
    
    trecIDList = trecIDDict.keys()
    trecIDList.sort(cmp=None, key=None, reverse=False)
    
    # The documents are anyway loading into a dict, so no worry about the sorting
    outputFileName = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/MLRelated/Model17/TOPK_sorted_by_queryID/distributedWorkForFeatureGenerationPostingRankInDoc20130727/" + inputRelativeFileName + "_priorityDocuments20130727"
    # outputFileName = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/MLRelated/Model17/TOPK_sorted_by_queryID/distributedWorkForFeatureGenerationPostingRankInDoc20130727/priorityDocuments20130727"
    outputFileHandler = open(outputFileName,"w")
    for trecID in trecIDList:
        outputFileHandler.write(trecID + "\n")
    outputFileHandler.close()
    
    # trecID example: GX000-71-13670316
    basePath = "/data/jhe/trecdata/"
    for trecID in trecIDList:
        trecIDElements = trecID.strip().split("-")
        completedPath = basePath + trecIDElements[0] + "/" + trecIDElements[1] + ".gz"
        # print "completedPath:",completedPath
        if completedPath not in correspondingGov2CompressedFilesCompletedPathsList:
            correspondingGov2CompressedFilesCompletedPathsList.append(completedPath)
    
    # the sorting doesn't matter here, just for easy observation
    correspondingGov2CompressedFilesCompletedPathsList.sort(cmp=None, key=None, reverse=False)
    outputFileName2 = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/MLRelated/Model17/TOPK_sorted_by_queryID/distributedWorkForFeatureGenerationPostingRankInDoc20130727/" + inputRelativeFileName + "_gov2_files_head_sth_20130727"
    # outputFileName2 = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/MLRelated/Model17/TOPK_sorted_by_queryID/distributedWorkForFeatureGenerationPostingRankInDoc20130727/missingTOPKResults_gov2_files_head_sth_20130722"
    outputFileHandler2 = open(outputFileName2,"w")
    outputFileHandler2.write("21" + "\n")
    for path in correspondingGov2CompressedFilesCompletedPathsList:
        outputFileHandler2.write(path + "\n")
    outputFileHandler2.close()
    
    print "Overall Processing Statistics:"
    print "len(trecIDList):",len(trecIDList)
    print "len(correspondingGov2CompressedFilesCompletedPathsList):",len(correspondingGov2CompressedFilesCompletedPathsList)
    print "inputFileName:",inputFileName
    print "outputFileName:",outputFileName
    print "outputFileName2:",outputFileName2


# currently do NOT use this anymore since 2013/07/14 night
def postprocessNaive():
    # 1st level:
    # key: trecID
    # value: a list containing many tuples 
        # tuple format:
        # (term,partialBM25Score,postingRankInDoc)
    trecIDWithPostingsListDict = {}
    
    # 1st level:
    # key: trecID
    # value: another dict
        # 2ed level:
        # key: term
        # value: partial bm25score
    trecIDWithPostingsDictDict = {}
    
    print "Open the file for file correctness check"
    inputFileName1 = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/gov2DocumentWithTheirCompletedPostingSet_for_debug_ONLY"
    inputFileHandler1 = open(inputFileName1,"r")
    inputFileHandler1.readline()
    for index,line in enumerate(inputFileHandler1.readlines()):
        lineElements = line.strip().split(" ")
        trecID = lineElements[0]
        # step1: check whether the docSizeInWords has been computed correctly.
        value1InInt = int( lineElements[1] )
        value2InInt = int( lineElements[2] )
        if value1InInt == value2InInt:
            pass
        else:
            print "mark1:",value1InInt,value2InInt
        
        # step2: check whether we have enough postings actually recorded
        if value1InInt == len(lineElements[3:]):
            pass
        else:
            print "mark2:",value1InInt,len(lineElements[3:])
        
        print index,trecID,len(lineElements)
    print "Check Result: The file correctness seems to be OK"
    inputFileHandler1.close()
    
    print "Open the file again for loading the postings for each document"
    inputFileName1 = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/gov2DocumentWithTheirCompletedPostingSet_for_debug_ONLY"
    inputFileHandler1 = open(inputFileName1,"r")
    inputFileHandler1.readline()
    for index,line in enumerate(inputFileHandler1.readlines()):
        lineElements = line.strip().split(" ")
        trecID = lineElements[0]
        if trecID not in trecIDWithPostingsListDict:
            trecIDWithPostingsListDict[trecID] = []
        
        for postingInString in lineElements[3:]:
            postingElements = postingInString.strip().split("(")[1].strip().split(")")[0].split(",")
            term = postingElements[0]
            scoreInFloatFormat = float( postingElements[1] )
            infoTuple = (term,scoreInFloatFormat,0) # init the posting rank in the doc to be 0
            trecIDWithPostingsListDict[trecID].append(infoTuple)
        
        # need to sort the trecIDWithPostingsListDict[trecID] based on the partialBM25 score
        trecIDWithPostingsListDict[trecID].sort(cmp=None, key=itemgetter(1), reverse=True)
        
        # need to compute the feature: posting rank in the doc
        previousTerm = ""
        previousScoreInFloatFormat = 0.0
        previousPostingRankInDoc = -1
        for index,infoTuple in enumerate(trecIDWithPostingsListDict[trecID]):
            oldTuple = infoTuple
            (currentTerm,currentScoreInFloatFormat,currentPostingRankInDoc) = oldTuple
            # if they have the same score, they should have the same rank
            if currentScoreInFloatFormat == previousScoreInFloatFormat:
                currentPostingRankInDoc = previousPostingRankInDoc
            # if they don't have the same score, then their ranks should be different
            else:
                currentPostingRankInDoc = index + 1
            newTuple = (currentTerm,currentScoreInFloatFormat,currentPostingRankInDoc)
            trecIDWithPostingsListDict[trecID][index] = newTuple
            
            
            # Update the data structure: trecIDWithPostingsDictDict
            if trecID not in trecIDWithPostingsDictDict:
                trecIDWithPostingsDictDict[trecID] = {}
                if currentTerm not in trecIDWithPostingsDictDict[trecID]:
                    trecIDWithPostingsDictDict[trecID][currentTerm] = currentPostingRankInDoc
                else:
                    # a check for the correctness of the ranking
                    if trecIDWithPostingsDictDict[trecID][currentTerm] != currentPostingRankInDoc:
                        print 
                        print "system logic error, mark1"
                        print "error log begins..."
                        print "trecID:",trecID
                        print "currentTerm:",currentTerm
                        print "trecIDWithPostingsDictDict[trecID][currentTerm]:",trecIDWithPostingsDictDict[trecID][currentTerm]
                        print "currentPostingRankInDoc:",currentPostingRankInDoc
                        print "error log ends." 
                        exit(1)
            else:
                if currentTerm not in trecIDWithPostingsDictDict[trecID]:
                    trecIDWithPostingsDictDict[trecID][currentTerm] = currentPostingRankInDoc
                else:
                    # a check for the correctness of the ranking
                    if trecIDWithPostingsDictDict[trecID][currentTerm] != currentPostingRankInDoc:
                        print
                        print "system logic error, mark2"
                        print "error log begins..."
                        print "trecID:",trecID
                        print "currentTerm:",currentTerm
                        print "trecIDWithPostingsDictDict[trecID][currentTerm]:",trecIDWithPostingsDictDict[trecID][currentTerm]
                        print "currentPostingRankInDoc:",currentPostingRankInDoc
                        print "error log ends."
                        exit(1)
            
                      
            # assignment for the variable previousTerm, previousScoreInFloatFormat and previousPostingRankInDoc
            previousTerm = currentTerm
            previousScoreInFloatFormat =  currentScoreInFloatFormat
            previousPostingRankInDoc = currentPostingRankInDoc
        print "trecID:",trecID,len(trecIDWithPostingsListDict[trecID]),len(trecIDWithPostingsDictDict[trecID])
        
        
        if trecID == "GX246-48-15656237" or trecID == "GX249-49-11480489":
            print "debug begins..."
            print "trecID:",trecID
            print trecIDWithPostingsListDict[trecID]
            '''
            debugOutputFileName = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/outputFileForDebug"
            debugOutputFileHandler = open(debugOutputFileName,"a")
            # print trecIDWithPostingsListDict[trecID]
            debugOutputFileHandler.write(trecID + "\n")
            for tuple in trecIDWithPostingsListDict[trecID]:
                debugOutputFileHandler.write(tuple + "\n")
            debugOutputFileHandler.write("\n")
            debugOutputFileHandler.close()
            '''
            print "debug ends."
    
    print "debug begins..."
    print "len(trecIDWithPostingsListDict):",len(trecIDWithPostingsListDict)
    print "len(trecIDWithPostingsDictDict):",len(trecIDWithPostingsDictDict)
    print "trecIDWithPostingsDictDict['GX246-48-15656237']['23218']:",trecIDWithPostingsDictDict['GX246-48-15656237']['23218']
    print "trecIDWithPostingsDictDict['GX246-48-15656237']['of']:",trecIDWithPostingsDictDict['GX246-48-15656237']['of']
    
    print "trecIDWithPostingsDictDict['GX249-49-11480489']['dig0024']:",trecIDWithPostingsDictDict['GX249-49-11480489']['dig0024']
    print "trecIDWithPostingsDictDict['GX249-49-11480489']['the']:",trecIDWithPostingsDictDict['GX249-49-11480489']['the']
    print "debug ends."
     
    inputFileHandler1.close()

def postprocessRational(inputFileNameForProvidingPostingRankInDoc,inputFileNameForPreviousTrainingFile,outputFileNameForNewTrainingFile):
    print "postprocessRational(...) called."
    outputFileHandler = open(outputFileNameForNewTrainingFile,"w")
    
    inputFileHandler1 = open(inputFileNameForProvidingPostingRankInDoc,"r")
    # SWITCH for ON/OFF for headline
    # inputFileHandler1.readline()
    
    currentLineForFile1 = ""
    currentLineForFile2 = ""
    
    currentTrecIDForFile1 = ""
    currentTrecIDForFile2 = ""
    
    currentTrecIDPostingList = []
    currentTrecIDTermANDRankDict = {}
    
    inputFileHandler2 = open(inputFileNameForPreviousTrainingFile,"r")
    
    # read the first line of the training example
    currentLineForFile2 = inputFileHandler2.readline()
    ###############################################
    while currentLineForFile2:
        currentLineForFile2Elements = currentLineForFile2.strip().split(" ")
        currentTrecIDForFile2 = currentLineForFile2Elements[2]
        currentTermForFile2 = currentLineForFile2Elements[4]
        
        # step1: loading process if necessary
        if currentTrecIDForFile2 != currentTrecIDForFile1:
            currentLineForFile1 = inputFileHandler1.readline()
            currentLineForFile1Elements = currentLineForFile1.strip().split(" ")
            currentTrecIDForFile1 = currentLineForFile1Elements[0]
            
            if currentTrecIDForFile1 != currentTrecIDForFile2:
                print "system error"
                print "currentTrecIDForFile1:",currentTrecIDForFile1
                print "currentTrecIDForFile2:",currentTrecIDForFile2
                exit(1)
            else:
                # in this case: currentTrecIDForFile1 == currentTrecIDForFile2
                # empty the data structure here in order to put in the new data
                currentTrecIDPostingList = []
                currentTrecIDTermANDRankDict = {}
                
                for postingInString in currentLineForFile1Elements[3:]:
                    postingElements = postingInString.strip().split("(")[1].strip().split(")")[0].split(",")
                    term = postingElements[0]
                    scoreInFloatFormat = float( postingElements[1] )
                    infoTuple = (term,scoreInFloatFormat,0) # init the posting rank in the doc to be 0
                    currentTrecIDPostingList.append(infoTuple)
                
                # need to sort the trecIDWithPostingsListDict[trecID] based on the partialBM25 score
                currentTrecIDPostingList.sort(cmp=None, key=itemgetter(1), reverse=True)
                
                # need to compute the feature: posting rank in the doc
                previousTerm = ""
                previousScoreInFloatFormat = 0.0
                previousPostingRankInDoc = -1
                for index,infoTuple in enumerate(currentTrecIDPostingList):
                    oldTuple = infoTuple
                    (currentTerm,currentScoreInFloatFormat,currentPostingRankInDoc) = oldTuple
                    # if they have the same score, they should have the same rank
                    if currentScoreInFloatFormat == previousScoreInFloatFormat:
                        currentPostingRankInDoc = previousPostingRankInDoc
                    # if they don't have the same score, then their ranks should be different
                    else:
                        currentPostingRankInDoc = index + 1
                    newTuple = (currentTerm,currentScoreInFloatFormat,currentPostingRankInDoc)
                    currentTrecIDPostingList[index] = newTuple
                    
                    if currentTerm not in currentTrecIDTermANDRankDict:
                        currentTrecIDTermANDRankDict[currentTerm] = currentPostingRankInDoc
                    else:
                        # a check for the correctness of the ranking
                        if currentTrecIDTermANDRankDict[currentTerm] != currentPostingRankInDoc:
                            print
                            print "system logic error, mark2"
                            print "error log begins..."
                            print "currentTrecIDForFile1:",currentTrecIDForFile1
                            print "currentTerm:",currentTerm
                            print "currentTrecIDTermANDRankDict[currentTerm]:",currentTrecIDTermANDRankDict[currentTerm]
                            print "currentPostingRankInDoc:",currentPostingRankInDoc
                            print "error log ends."
                            exit(1)
                             
                    # assignment for the variable previousTerm, previousScoreInFloatFormat and previousPostingRankInDoc
                    previousTerm = currentTerm
                    previousScoreInFloatFormat =  currentScoreInFloatFormat
                    previousPostingRankInDoc = currentPostingRankInDoc
                    
                print "Loaded",currentTrecIDForFile1,len(currentTrecIDPostingList),len(currentTrecIDTermANDRankDict)
        else:
            pass
            # print "data from file1 can handle the request from file2"
        
        '''
        # for debug ONLY
        if currentTrecIDForFile1 == "GX249-49-11480489" or currentTrecIDForFile1 == "GX257-13-0667777" or currentTrecIDForFile1 == "GX231-92-13503936":
            print "currentTrecIDForFile1:",currentTrecIDForFile1
            print "currentTrecIDPostingList:",currentTrecIDPostingList
            print "currentTrecIDTermANDRankDict:",currentTrecIDTermANDRankDict
        '''
        
        # step2: assigning new feature value process
        valueForPostingRankInDoc = currentTrecIDTermANDRankDict[currentTermForFile2]
        outputFileHandler.write(currentLineForFile2.strip() + " " + str(valueForPostingRankInDoc) + "\n")
        
        currentLineForFile2 = inputFileHandler2.readline()
    ###############################################
    outputFileHandler.close()
    inputFileHandler1.close()
    inputFileHandler2.close()
    
def checkInputFile(inputFileName1):
    print "checkInputFile(...) called."
    print "Purpose: Open the file for checking the file correctness"
    inputFileHandler1 = open(inputFileName1,"r")
    
    # SWITCH for ON/OFF for the headline
    inputFileHandler1.readline()
    
    fileLines = inputFileHandler1.readlines()
    for index,line in enumerate(fileLines):
        lineElements = line.strip().split(" ")
        trecID = lineElements[0]
        # step1: check whether the docSizeInWords has been computed correctly.
        value1InInt = int( lineElements[1] )
        value2InInt = int( lineElements[2] )
        if value1InInt == value2InInt:
            pass
        else:
            print "mark1:",value1InInt,value2InInt
        
        # step2: check whether we have enough postings actually recorded
        if value1InInt == len(lineElements[3:]):
            pass
        else:
            print "mark2:",value1InInt,len(lineElements[3:])
        
        print index,trecID,len(lineElements)
    print "Check Result: The file correctness seems to be OK"
    
    #######################
    # in this case: currentTrecIDForFile1 == currentTrecIDForFile2
    # empty the data structure here in order to put in the new data
    currentLineForFile1Elements = fileLines[0].strip().split(" ")
    currentTrecIDForFile1 = currentLineForFile1Elements[0]
    
    currentTrecIDPostingList = []
    currentTrecIDTermANDRankDict = {}
    
    for postingInString in currentLineForFile1Elements[3:]:
        postingElements = postingInString.strip().split("(")[1].strip().split(")")[0].split(",")
        term = postingElements[0]
        scoreInFloatFormat = float( postingElements[1] )
        infoTuple = (term,scoreInFloatFormat,0) # init the posting rank in the doc to be 0
        currentTrecIDPostingList.append(infoTuple)
    
    # need to sort the trecIDWithPostingsListDict[trecID] based on the partialBM25 score
    currentTrecIDPostingList.sort(cmp=None, key=itemgetter(1), reverse=True)
    
    # need to compute the feature: posting rank in the doc
    previousTerm = ""
    previousScoreInFloatFormat = 0.0
    previousPostingRankInDoc = -1
    for index,infoTuple in enumerate(currentTrecIDPostingList):
        oldTuple = infoTuple
        (currentTerm,currentScoreInFloatFormat,currentPostingRankInDoc) = oldTuple
        # if they have the same score, they should have the same rank
        if currentScoreInFloatFormat == previousScoreInFloatFormat:
            currentPostingRankInDoc = previousPostingRankInDoc
        # if they don't have the same score, then their ranks should be different
        else:
            currentPostingRankInDoc = index + 1
        newTuple = (currentTerm,currentScoreInFloatFormat,currentPostingRankInDoc)
        currentTrecIDPostingList[index] = newTuple
        
        if currentTerm not in currentTrecIDTermANDRankDict:
            currentTrecIDTermANDRankDict[currentTerm] = currentPostingRankInDoc
        else:
            # a check for the correctness of the ranking
            if currentTrecIDTermANDRankDict[currentTerm] != currentPostingRankInDoc:
                print
                print "system logic error, mark2"
                print "error log begins..."
                print "currentTrecIDForFile1:",currentTrecIDForFile1
                print "currentTerm:",currentTerm
                print "currentTrecIDTermANDRankDict[currentTerm]:",currentTrecIDTermANDRankDict[currentTerm]
                print "currentPostingRankInDoc:",currentPostingRankInDoc
                print "error log ends."
                exit(1)
                 
        # assignment for the variable previousTerm, previousScoreInFloatFormat and previousPostingRankInDoc
        previousTerm = currentTerm
        previousScoreInFloatFormat =  currentScoreInFloatFormat
        previousPostingRankInDoc = currentPostingRankInDoc
        
    print currentTrecIDForFile1,len(currentTrecIDPostingList),len(currentTrecIDTermANDRankDict),"Loaded"
    # for debug ONLY
    # print "currentTrecIDPostingList:",currentTrecIDPostingList
    # print "currentTrecIDTermANDRankDict:",currentTrecIDTermANDRankDict
    # for postingTuple in currentTrecIDPostingList:
    #    print postingTuple
    # print "currentTrecIDTermANDRankDict:",currentTrecIDTermANDRankDict
    #######################
    
    inputFileHandler1.close()

print "Program Begins..."
# This function is used for generating the inputs for the polyIRToolkit
# The following are the outputs of the function preprocess()
# (1) priorityDocumentsNeededForFeatureGeneration
# (2) gov2_files_head_sth
# preprocess()

# input files
inputFileNameForProvidingPostingRankInDoc = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/MLRelated/Model17/TOPK_sorted_by_queryID/distributedWorkForFeatureGenerationPostingRankInDoc20130727/gov2DocumentWithTheirCompletedPostingSet_for_debug_ONLY_xaj_sorted_by_trecID"
inputFileNameForPreviousTrainingFile = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/MLRelated/Model17/TOPK_sorted_by_queryID/distributedWorkForFeatureGenerationPostingRankInDoc20130727/xaj"
outputFileNameForNewTrainingFile = "/data3/obukai/the_new_trip_of_feature_generation/gov2ClearYourMindAndDoItAgain/MLRelated/Model17/TOPK_sorted_by_queryID/distributedWorkForFeatureGenerationPostingRankInDoc20130727/xaj_postingRankInDocAdded_WITHOUT_HEADLINE.train"

# This function is used for checking the completeness of the output file of the polyIRToolkit
# The output file is somewhat like: gov2DocumentWithTheirCompletedPostingSet_for_debug_ONLY_xaa
# checkInputFile(inputFileNameForProvidingPostingRankInDoc)

# current used version
# The file inputFileNameForProvidingPostingRankInDoc needs to be sorted as the same as the inputFileNameForPreviousTrainingFile
postprocessRational(inputFileNameForProvidingPostingRankInDoc,inputFileNameForPreviousTrainingFile,outputFileNameForNewTrainingFile)
print "Program Ends."


